import sys

sys.path.append('/data/jinzhuoran/cognlp')
sys.path.append('/data/zhuoran/cognlp')
sys.path.append('/data/zhuoran/code/cognlp')

import torch
import torch.nn as nn
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler
from torch.utils.data import DataLoader, RandomSampler
from cognlp import *
from cognlp.io.loader.rex import TrexReLoader
from cognlp.io.processor.rex import REXProcessor
from cognlp.models.re.bert_re import Bert4Re, Bert4ReParallel
from cognlp.core.metrics import AccuracyMetric
from cognlp.core.trainer import Trainer, Tester
from cognlp.core.dataset import RexDataset
from cognlp.utils.util import load_json

torch.cuda.set_device(0)

device = torch.device('cuda')
vocabulary = Vocabulary.load('../data/re/trexfull/data/vocabulary.txt')
data = DataTable.load_table(path='../data/re/trexfull/data/train.json')
train_data, dev_data, test_data = data.split(9, 1, 0)


train_data = RexDataset(train_data)
train_sampler = RandomSampler(train_data)
dev_data = RexDataset(dev_data)
dev_sampler = RandomSampler(dev_data)
model = Bert4Re(len(vocabulary))

loss = nn.CrossEntropyLoss(ignore_index=0)
optimizer = optim.Adam(model.parameters(), lr=0.00005)
metric = ClassifyFPreRecMetric()

trainer = Trainer(model,
                  train_data,
                  dev_data=dev_data,
                  n_epochs=20,
                  batch_size=100,
                  loss=loss,
                  optimizer=optimizer,
                  scheduler=None,
                  metrics=metric,
                  train_sampler=train_sampler,
                  dev_sampler=dev_sampler,
                  drop_last=False,
                  gradient_accumulation_steps=1,
                  num_workers=0,
                  save_path="../data/re/trexfull/model",
                  save_file=None,
                  print_every=None,
                  scheduler_steps=None,
                  validate_steps=None,
                  save_steps=None,
                  grad_norm=None,
                  use_tqdm=True,
                  device=device,
                  device_ids=[0, 3, 4, 5],
                  callbacks=None,
                  metric_key=None,
                  writer_path='../data/re/trexfull/tensorboard',
                  fp16=False,
                  fp16_opt_level='O1',
                  seed=527,
                  checkpoint_path=None,
                  task='ace2005-new',
                  logger_path='../data/re/trexfull/logger')

trainer.train()
print(1)
